{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"RERUN_NOTEBOOK_ASSET\"] = \"serve-local\"\n",
    "\n",
    "import os\n",
    "\n",
    "import cv2\n",
    "import pyarrow as pa\n",
    "import requests\n",
    "import rerun as rr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download(url: str) -> None:\n",
    "    if not os.path.exists(\"datasets\"):\n",
    "        os.makedirs(\"datasets\")\n",
    "    filename = url.split(\"/\")[-1]\n",
    "    dest = f\"datasets/{filename}\"\n",
    "    if not os.path.exists(dest):\n",
    "        with open(dest, \"wb\") as f:\n",
    "            f.write(requests.get(url).content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# OPF\n",
    "download(\"http://app.rerun.io/version/0.18.2/examples/open_photogrammetry_format.rrd\")\n",
    "opf = rr.dataframe.load_recording(\"datasets/open_photogrammetry_format.rrd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What timelines do we have?\n",
    "opf.schema().index_columns()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What components do we have?\n",
    "opf.schema().component_columns()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract a single point-cloud\n",
    "\n",
    "cloud = rr.dataframe.ComponentColumnSelector(\"world/points\", rr.components.Position3D)\n",
    "\n",
    "batches = opf.view(index=\"image\", contents=\"world/points\").select([cloud])\n",
    "\n",
    "# Should the python APIs do this for us?\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "cloud = table[0][0]\n",
    "type(cloud)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "points_np = cloud.values.flatten().to_numpy().reshape(-1, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "points_np[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the positions of all the camera estimates\n",
    "\n",
    "img_num = rr.dataframe.IndexColumnSelector(\"image\")\n",
    "logged = rr.dataframe.IndexColumnSelector(\"log_time\")\n",
    "\n",
    "pos = rr.dataframe.ComponentColumnSelector(\"world/cameras\", rr.components.Translation3D.component_type())\n",
    "quat = rr.dataframe.ComponentColumnSelector(\"world/cameras\", rr.components.TransformMat3x3.component_type())\n",
    "\n",
    "view = opf.view(index=\"image\", contents=\"world/cameras\")\n",
    "batches = view.select(img_num, logged, pos, quat)\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Annoying -- row traversal even though they were logged together :-(\n",
    "table = table.combine_chunks()\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at some samples (Polars would make this nicer)\n",
    "table.take([0, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Let's feed this into a geometry library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Maybe we only wanted to query poses for a few images\n",
    "# Find the positions of all the camera estimates\n",
    "\n",
    "batches = view.filter_range_sequence(start=10, end=20).select([img_num, logged, pos, quat])\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "table.take([0, 10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: latest-at"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's grab an image\n",
    "\n",
    "view = opf.view(index=\"image\", contents=\"/**\").filter_range_sequence(42, 42)\n",
    "\n",
    "img_num = rr.dataframe.IndexColumnSelector(\"image\")\n",
    "\n",
    "# Getting these right is annoying -- would be nice to support '/**' somehow\n",
    "blob = rr.dataframe.ComponentColumnSelector(\"world/cameras/image\", rr.components.Blob.component_type())\n",
    "media = rr.dataframe.ComponentColumnSelector(\"world/cameras/image\", rr.components.MediaType.component_type())\n",
    "\n",
    "batches = view.select(blob, media)\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ^^ when we typo something, we just find no data\n",
    "table.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's grab an image for real this time\n",
    "\n",
    "# Getting these right is annoying -- would be nice to support '/**' somehow\n",
    "blob = rr.dataframe.ComponentColumnSelector(\"world/cameras/image/rgb\", rr.components.Blob)\n",
    "media = rr.dataframe.ComponentColumnSelector(\"world/cameras/image/rgb\", rr.components.MediaType)\n",
    "\n",
    "batches = view.select([blob, media])\n",
    "\n",
    "table = pa.Table.from_batches(batches)\n",
    "\n",
    "print(f\"Found {table.num_rows} rows in {len(table.columns[0].chunks)} chunks\")\n",
    "\n",
    "## NOTE: Kind of surprising to find 2 rows here. The first one gets included because the VIEW contains\n",
    "## Contains static data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Image is encoded as: {table[1][1][0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = image = cv2.imdecode(table[0][1][0].values.to_numpy(), cv2.IMREAD_COLOR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"preview\")\n",
    "\n",
    "rr.log(\"img\", rr.Image(image, color_model=\"bgr\"))\n",
    "\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# BGR shader not working on web in firefox :sob:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"preview\")\n",
    "\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "rr.log(\"img\", rr.Image(image))\n",
    "\n",
    "rr.notebook_show()"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
